package com.yc154.web;

import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 登录过滤器: 拦截登录请求, 如果登录成功, 则生成 JWT 令牌, 通过响应发送给浏览器
 */
@Component
public class LoginFilterGatewayFilterFactory extends
        AbstractGatewayFilterFactory<Object> implements Ordered {

    @Override
    public int getOrder() {
        return 0;
    }

    @Override
    public String name() {
        return "loginFilter";
    }

    @Override
    public GatewayFilter apply(Object config) {
        return (exchange, chain) -> {
            ServerHttpResponse originalResponse = exchange.getResponse();
            BodyCaptureResponse bodyCaptureResponse = new BodyCaptureResponse(originalResponse);

            return chain.filter(exchange.mutate().response(bodyCaptureResponse).build())
                    .then(Mono.defer(() -> {
                        // 获取捕获的响应体
                        byte[] body = bodyCaptureResponse.getBody();
                        String resultStr = new String(body);
                        System.out.println("响应内容: " + resultStr);

                        // 只有在有响应体的情况下才处理
                        if (body.length > 0) {
                            // 添加 JWT 头
                            originalResponse.getHeaders().add("jwt", "testJwt");

                            // 重新写入响应体
                            DataBuffer buffer = originalResponse.bufferFactory().wrap(body);
                            return originalResponse.writeWith(Mono.just(buffer));
                        } else {
                            return originalResponse.setComplete();
                        }
                    }));
        };
    }
}

class BodyCaptureResponse extends ServerHttpResponseDecorator {
    private final ByteArrayOutputStream bodyStream = new ByteArrayOutputStream();
    private final AtomicReference<DataBuffer> joinedBuffer = new AtomicReference<>();

    public BodyCaptureResponse(ServerHttpResponse delegate) {
        super(delegate);
    }

    @Override
    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
        return Mono.from(body).flatMap(buffer -> {
            // 捕获数据
            byte[] bytes = new byte[buffer.readableByteCount()];
            buffer.read(bytes);
            try {
                bodyStream.write(bytes);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }

            // 重新包装数据供后续使用
            DataBuffer wrappedBuffer = getDelegate().bufferFactory().wrap(bytes);
            joinedBuffer.set(wrappedBuffer);

            return super.writeWith(Mono.just(wrappedBuffer));
        });
    }

    public byte[] getBody() {
        return bodyStream.toByteArray();
    }

    public DataBuffer getBodyBuffer() {
        return joinedBuffer.get();
    }
}